from typing import Dict, List, Any, Optional
import json
from loguru import logger

class AnswerGenerator:
    """
    回答生成器，负责处理子图数据和生成最终回答
    """
    def __init__(self, ollama_client):
        self.ollama_client = ollama_client
    
    async def generate_answer(self, question: str, subgraph: Dict[str, Any], intent: Dict[str, Any], temperature: float = 0.7) -> str:
        """
        根据子图数据生成回答
        
        Args:
            question: 用户问题
            subgraph: 子图数据
            intent: 意图解析结果
            temperature: 温度参数，控制回答的创造性
            
        Returns:
            生成的回答文本
        """
        # 将子图数据转换为结构化文本描述
        subgraph_text = self._format_subgraph(subgraph)
        
        # 构建提示词
        system_prompt = """
        你是一个基于知识图谱的问答助手。根据提供的子图信息，回答用户的问题。
        遵循以下规则：
        1. 只使用子图中包含的信息来回答问题
        2. 如果子图信息不足，可以基于常识进行补充，但需明确标注「这是基于我的知识补充的信息，不在知识图谱中」
        3. 回答应简洁、准确、通顺
        4. 使用用户提问的语言回答
        5. 如果子图为空，请告知用户未找到相关信息
        """
        
        prompt = f"问题: {question}\n\n子图信息:\n{subgraph_text}\n\n意图: {json.dumps(intent, ensure_ascii=False)}"
        
        # 调用LLM生成回答
        try:
            answer = await self.ollama_client.generate(prompt, system_prompt, temperature=temperature)
            return answer
        except Exception as e:
            logger.error(f"生成回答失败: {str(e)}")
            return f"抱歉，生成回答时出现错误: {str(e)}"
    
    def _format_subgraph(self, subgraph: Dict[str, Any]) -> str:
        """
        将子图数据格式化为文本描述
        
        Args:
            subgraph: 子图数据
            
        Returns:
            格式化的文本描述
        """
        if not subgraph or (not subgraph.get("nodes") and not subgraph.get("edges")):
            return "子图为空，未找到相关信息。"
        
        # 创建节点ID到节点数据的映射，方便后续引用
        node_map = {node['id']: node for node in subgraph.get("nodes", [])}
        
        # 格式化节点信息
        nodes_text = ""
        for node in subgraph.get("nodes", []):
            # 过滤掉空属性，并格式化属性值
            attributes = {k: v for k, v in node.get("attributes", {}).items() if v is not None and v != ""}
            attr_text = ", ".join([f"{k}: {v}" for k, v in attributes.items()])
            nodes_text += f"节点ID: {node['id']}, 标签: {node['label']}, 属性: {{{attr_text}}}\n"
        
        # 格式化边信息，并添加节点标签以增强可读性
        edges_text = ""
        for edge in subgraph.get("edges", []):
            source_node = node_map.get(edge['source'], {})
            target_node = node_map.get(edge['target'], {})
            source_label = source_node.get('label', edge['source'])
            target_label = target_node.get('label', edge['target'])
            
            # 过滤掉空属性
            attributes = {k: v for k, v in edge.get("attributes", {}).items() if v is not None and v != ""}
            attr_text = ", ".join([f"{k}: {v}" for k, v in attributes.items()])
            
            edges_text += f"关系: {source_label}({edge['source']}) --[{edge['type']}]--> {target_label}({edge['target']}), 属性: {{{attr_text}}}\n"
        
        # 添加子图摘要信息
        summary = f"子图包含 {len(subgraph.get('nodes', []))} 个节点和 {len(subgraph.get('edges', []))} 个关系。\n\n"
        
        return f"{summary}节点:\n{nodes_text}\n关系:\n{edges_text}"
    
    async def direct_answer(self, question: str, temperature: float = 0.7) -> str:
        """
        直接使用大语言模型回答问题（不使用知识图谱）
        
        Args:
            question: 用户问题
            temperature: 温度参数，控制回答的创造性
            
        Returns:
            模型生成的回答
        """
        system_prompt = """
        你是一个智能助手，请直接回答用户的问题。
        如果你不确定答案，请诚实地表明。
        使用用户提问的语言回答。
        """
        
        try:
            return await self.ollama_client.generate(question, system_prompt, temperature=temperature)
        except Exception as e:
            logger.error(f"直接回答失败: {str(e)}")
            return f"抱歉，生成回答时出现错误: {str(e)}"
    
    async def fallback_answer(self, question: str, error_message: str = None, temperature: float = 0.7) -> str:
        """
        当知识图谱查询失败时的备用回答
        
        Args:
            question: 用户问题
            error_message: 错误信息
            
        Returns:
            备用回答
        """
        system_prompt = """
        你是一个智能助手。用户的问题无法通过知识图谱回答，请提供以下回复：
        1. 告知用户无法在知识图谱中找到相关信息
        2. 如果可能，提供一个基于你自身知识的简短回答
        3. 建议用户如何重新表述问题以获得更好的结果
        使用用户提问的语言回答。
        """
        
        prompt = f"问题: {question}\n\n错误信息: {error_message if error_message else '未找到相关信息'}"
        
        try:
            return await self.ollama_client.generate(prompt, system_prompt, temperature=temperature)
        except Exception as e:
            logger.error(f"生成备用回答失败: {str(e)}")
            return "抱歉，我无法在知识图谱中找到相关信息，并且生成备用回答时出现错误。请尝试重新表述您的问题。"